# -*- coding:utf8 -*-
import os
import csv
import random

def split_train_test():
    img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/train'
    csv_train = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/smallDataSet/one/train224/offline/trainFile.csv'
    csv_test = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/smallDataSet/one/train224/offline/testFile.csv'
    csv_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/smallDataSet/one/train224/csvFile'
    f_train = open(csv_train,'wb')
    f_test = open(csv_test,'wb')
    train_writer = csv.writer(f_train)
    test_writer = csv.writer(f_test)

    for file in os.listdir(csv_folder):
        flag = 0
        path = csv_folder + '/' + file
        f = open(path,'rb')
        lines = f.readlines()
        random.shuffle(lines)
        num = len(lines)
        for line in lines:
            flag += 1
            name, label = line.split(',')
            if flag <= int(num * 3 / 4):
                new_line = [img_folder + '/' + name, int(label)]
                train_writer.writerow(new_line)
            elif flag >= int(num * 3 / 4):
                new_line = [img_folder + '/' + name, int(label)]
                test_writer.writerow(new_line)
    f_test.close()
    f_train.close()

    f_train = open(csv_train,'rb')
    f_test = open(csv_test,'rb')
    train_lines = f_train.readlines()
    test_lines = f_test.readlines()
    random.shuffle(test_lines)
    random.shuffle(train_lines)
    f_train = open(csv_train,'wb')
    f_test = open(csv_test,'wb')
    train_writer = csv.writer(f_train)
    test_writer = csv.writer(f_test)
    for line in test_lines:
        name, label = line.split(',')
        new_line = [name, int(label)]
        test_writer.writerow(new_line)
    for line in train_lines:
        name, label = line.split(',')
        new_line = [name, int(label)]
        train_writer.writerow(new_line)

if __name__ == "__main__":
    split_train_test()